package rocket

import (
	"context"

	"github.com/apache/rocketmq-client-go/v2"
	"github.com/apache/rocketmq-client-go/v2/consumer"
	"github.com/apache/rocketmq-client-go/v2/primitive"
)

type (
	Consume func(key string, body []byte)

	Consumer interface {
		Consume(consume Consume)
		// 开启
		Start()
		// 关闭
		Stop()
	}

	defaultConsumer struct {
		consumer rocketmq.PushConsumer
		topic    string
	}
)

func NewConsumer(cfg *ConsumerConfig) Consumer {
	pAddr, err := primitive.NewNamesrvAddr(cfg.Addr...)
	if err != nil {
		panic(err)
	}

	opts := make([]consumer.Option, 0, 5)
	opts = append(opts, consumer.WithNameServer(pAddr))
	opts = append(opts, consumer.WithConsumerModel(consumer.Clustering))
	opts = append(opts, consumer.WithGroupName(cfg.Topic))
	if cfg.Auth != nil {
		opts = append(opts, consumer.WithCredentials(primitive.Credentials{
			AccessKey:     cfg.Auth.AccessKey,
			SecretKey:     cfg.Auth.SecretKey,
			SecurityToken: cfg.Auth.SecurityToken,
		}))
	}

	c, err := rocketmq.NewPushConsumer(opts...)
	if err != nil {
		panic(err)
	}

	return &defaultConsumer{
		consumer: c,
		topic:    cfg.Topic,
	}
}

func (d *defaultConsumer) Start() {
	err := d.consumer.Start()
	if err != nil {
		panic(err)
	}
}

func (d *defaultConsumer) Stop() {
	_ = d.consumer.Shutdown()
}

func (d *defaultConsumer) Consume(consume Consume) {
	err := d.consumer.Subscribe(d.topic, consumer.MessageSelector{}, func(ctx context.Context, ext ...*primitive.MessageExt) (consumer.ConsumeResult, error) {
		for _, messageExt := range ext {
			consume(messageExt.GetTags(), messageExt.Body)
		}
		return consumer.ConsumeSuccess, nil
	})
	if err != nil {
		panic(err)
	}
}
